package com.boroborome.kibana;

import com.happy3w.toolkits.utils.StepAction;
import com.opencsv.CSVWriter;
import org.springframework.boot.json.BasicJsonParser;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Spliterators;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

public class LogDownloader {
    public void download(String kibanaUrl,
                         String indexPattern,
                         String condition,
                         String outputFile,
                         String cookies,
                         List<String> fields) throws IOException {
        ensureDirExist(new File(outputFile).getParent());
        FileWriter outputPrinter = new FileWriter(outputFile);
        CSVWriter csvWriter = new CSVWriter(outputPrinter);
        StepAction stepAction = new StepAction(10000, n -> System.out.println("Load log:" + n));

        Stream<String[]> records = queryAllRecords(kibanaUrl, indexPattern, condition, cookies, fields);
        records.peek($ -> stepAction.increase())
                .forEach(record -> csvWriter.writeNext(record));
        csvWriter.close();
    }

    private void ensureDirExist(String dirName) {
        File dir = new File(dirName);
        if (!dir.exists()) {
            dir.mkdirs();
        }
    }

    private Stream<String[]> queryAllRecords(String kibanaUrl,
                                             String indexPattern,
                                             String condition,
                                             String cookies,
                                             List<String> fields) {
        return StreamSupport.stream(new RecordStream(kibanaUrl, indexPattern, condition, cookies, fields), false);
    }

    private static class RecordStream extends Spliterators.AbstractSpliterator<String[]> {
        private String kibanaUrl;
        private String indexPattern;
        private Iterator preResultIt;
        private String cookies;
        private String scrollId;
        private List<String> fields;

        HttpClient client = HttpClient.newHttpClient();
        private BasicJsonParser jsonParser = new BasicJsonParser();

        public RecordStream(String kibanaUrl, String indexPattern, String condition, String cookies, List<String> fields) {
            super(0, 0);
            this.kibanaUrl = kibanaUrl;
            this.indexPattern = indexPattern;
            this.cookies = cookies;
            this.fields = fields;

            List result = queryData(MessageFormat.format("{0}/api/console/proxy?path=%2F{1}%2F_search%3Fscroll%3D1m&method=POST",
                    kibanaUrl, indexPattern),
                    condition);

            this.preResultIt = result.iterator();
        }

        @Override
        public boolean tryAdvance(Consumer<? super String[]> action) {
            if (!preResultIt.hasNext()) {
                List result = queryData(MessageFormat.format("{0}/api/console/proxy?path=%2F_search%2Fscroll&method=POST",
                        kibanaUrl),
                        String.format("{\"scroll\": \"1m\", \"scroll_id\": \"%s\"}", scrollId));
                preResultIt = result.iterator();
                if (!preResultIt.hasNext()) {
                    return false;
                }
            }

            Map hit = getNextMap(preResultIt.next());
            Map source = (Map) hit.get("_source");

            String[] values = fields.stream()
                    .map(field -> (String) source.get(field))
                    .collect(Collectors.toList())
                    .toArray(new String[0]);

            action.accept(values);
            return true;
        }

        private Map getNextMap(Object next) {
            if (next instanceof Map) {
                return (Map) next;
            } else if (next instanceof String) {
                return jsonParser.parseMap(((String)next).trim());
            }
            throw new UnsupportedOperationException("Unknown type:" + next.getClass());
        }


        private List queryData(String url, String body) {
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(url))
                    .header("Content-Type", "application/json")
                    .header("kbn-version", "7.13.4")
                    .header("Cookie", cookies)
                    .POST(HttpRequest.BodyPublishers.ofString(body))
                    .build();
            try {
                HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
                if (response.statusCode() >= 400) {
                    return Collections.emptyList();
                }

                Map<String, Object> responseMap =jsonParser.parseMap(response.body());
                Object newScrollId = responseMap.get("_scroll_id");
                if (newScrollId != null) {
                    scrollId = String.valueOf(newScrollId);
                }
                return (List) ((Map) responseMap.get("hits")).get("hits");

            } catch (Exception e) {
                throw new RuntimeException(e.getMessage(), e);
            }
        }
    }
}
